// sm2_demo.c
#include <stdio.h>
#include <string.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/err.h>
#include <openssl/core_names.h>

// 错误打印辅助函数
void print_errors(void) {
    ERR_print_errors_fp(stderr);
}

// 生成 SM2 密钥对并保存为 PEM 文件
int generate_sm2_keypair(const char *priv_key_file, const char *pub_key_file) {
    EVP_PKEY *pkey = NULL;
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_SM2, NULL);
    if (!ctx) {
        fprintf(stderr, "EVP_PKEY_CTX_new_id failed\n");
        print_errors();
        return -1;
    }

    if (EVP_PKEY_keygen_init(ctx) <= 0) {
        fprintf(stderr, "EVP_PKEY_keygen_init failed\n");
        print_errors();
        EVP_PKEY_CTX_free(ctx);
        return -1;
    }

    // 设置 SM2 曲线（可选，OpenSSL 3.0 中 SM2 默认使用 SM2 曲线）
    if (EVP_PKEY_CTX_set_group_name(ctx, SN_sm2) <= 0) {
        fprintf(stderr, "EVP_PKEY_CTX_set_group_name failed\n");
        print_errors();
        EVP_PKEY_CTX_free(ctx);
        return -1;
    }

    if (EVP_PKEY_generate(ctx, &pkey) <= 0) {
        fprintf(stderr, "EVP_PKEY_generate failed\n");
        print_errors();
        EVP_PKEY_CTX_free(ctx);
        return -1;
    }
    EVP_PKEY_CTX_free(ctx);

    // 保存私钥
    FILE *fpriv = fopen(priv_key_file, "w");
    if (!fpriv || !PEM_write_PrivateKey(fpriv, pkey, NULL, NULL, 0, NULL, NULL)) {
        fprintf(stderr, "Failed to write private key\n");
        print_errors();
        fclose(fpriv);
        EVP_PKEY_free(pkey);
        return -1;
    }
    fclose(fpriv);

    // 保存公钥
    FILE *fpub = fopen(pub_key_file, "w");
    if (!fpub || !PEM_write_PUBKEY(fpub, pkey)) {
        fprintf(stderr, "Failed to write public key\n");
        print_errors();
        fclose(fpub);
        EVP_PKEY_free(pkey);
        return -1;
    }
    fclose(fpub);

    printf("SM2 keypair generated: %s (private), %s (public)\n", priv_key_file, pub_key_file);
    EVP_PKEY_free(pkey);
    return 0;
}

// SM2 加密
int sm2_encrypt(EVP_PKEY *pub_key, const unsigned char *plaintext, size_t pt_len,
                unsigned char **ciphertext, size_t *ct_len) {
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(pub_key, NULL);
    if (!ctx) return -1;

    if (EVP_PKEY_encrypt_init(ctx) <= 0) goto err;
    if (EVP_PKEY_CTX_set_rsa_padding(ctx, RSA_PKCS1_PADDING) <= 0) {
        // 注意：SM2 不使用 RSA padding，但某些 OpenSSL 版本需要显式设置 SM2 模式
        // 实际上，OpenSSL 3.0+ 会自动识别 SM2
    }

    // 获取密文长度
    if (EVP_PKEY_encrypt(ctx, NULL, ct_len, plaintext, pt_len) <= 0) goto err;

    *ciphertext = OPENSSL_malloc(*ct_len);
    if (!*ciphertext) goto err;

    if (EVP_PKEY_encrypt(ctx, *ciphertext, ct_len, plaintext, pt_len) <= 0) {
        OPENSSL_free(*ciphertext);
        *ciphertext = NULL;
        goto err;
    }

    EVP_PKEY_CTX_free(ctx);
    return 0;

err:
    print_errors();
    EVP_PKEY_CTX_free(ctx);
    return -1;
}

// SM2 解密
int sm2_decrypt(EVP_PKEY *priv_key, const unsigned char *ciphertext, size_t ct_len,
                unsigned char **plaintext, size_t *pt_len) {
    EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new(priv_key, NULL);
    if (!ctx) return -1;

    if (EVP_PKEY_decrypt_init(ctx) <= 0) goto err;

    if (EVP_PKEY_decrypt(ctx, NULL, pt_len, ciphertext, ct_len) <= 0) goto err;

    *plaintext = OPENSSL_malloc(*pt_len);
    if (!*plaintext) goto err;

    if (EVP_PKEY_decrypt(ctx, *plaintext, pt_len, ciphertext, ct_len) <= 0) {
        OPENSSL_free(*plaintext);
        *plaintext = NULL;
        goto err;
    }

    EVP_PKEY_CTX_free(ctx);
    return 0;

err:
    print_errors();
    EVP_PKEY_CTX_free(ctx);
    return -1;
}

// SM2 签名（使用 SM3 作为摘要）
int sm2_sign(EVP_PKEY *priv_key, const unsigned char *data, size_t data_len,
             unsigned char **sig, size_t *sig_len) {
    EVP_MD_CTX *md_ctx = EVP_MD_CTX_new();
    if (!md_ctx) return -1;

    if (EVP_DigestSignInit(md_ctx, NULL, EVP_sm3(), NULL, priv_key) <= 0) goto err;
    if (EVP_DigestSignUpdate(md_ctx, data, data_len) <= 0) goto err;
    if (EVP_DigestSignFinal(md_ctx, NULL, sig_len) <= 0) goto err;

    *sig = OPENSSL_malloc(*sig_len);
    if (!*sig) goto err;

    if (EVP_DigestSignFinal(md_ctx, *sig, sig_len) <= 0) {
        OPENSSL_free(*sig);
        *sig = NULL;
        goto err;
    }

    EVP_MD_CTX_free(md_ctx);
    return 0;

err:
    print_errors();
    EVP_MD_CTX_free(md_ctx);
    return -1;
}

// SM2 验签
int sm2_verify(EVP_PKEY *pub_key, const unsigned char *data, size_t data_len,
               const unsigned char *sig, size_t sig_len) {
    EVP_MD_CTX *md_ctx = EVP_MD_CTX_new();
    if (!md_ctx) return -1;

    if (EVP_DigestVerifyInit(md_ctx, NULL, EVP_sm3(), NULL, pub_key) <= 0) goto err;
    if (EVP_DigestVerifyUpdate(md_ctx, data, data_len) <= 0) goto err;
    int ret = EVP_DigestVerifyFinal(md_ctx, sig, sig_len);

    EVP_MD_CTX_free(md_ctx);
    if (ret != 1) {
        print_errors();
        return -1;
    }
    return 0;

err:
    print_errors();
    EVP_MD_CTX_free(md_ctx);
    return -1;
}

// 加载私钥
EVP_PKEY *load_private_key(const char *filename) {
    FILE *f = fopen(filename, "r");
    if (!f) {
        perror("fopen private key");
        return NULL;
    }
    EVP_PKEY *pkey = PEM_read_PrivateKey(f, NULL, NULL, NULL);
    fclose(f);
    if (!pkey) {
        fprintf(stderr, "Failed to load private key from %s\n", filename);
        print_errors();
    }
    return pkey;
}

// 加载公钥
EVP_PKEY *load_public_key(const char *filename) {
    FILE *f = fopen(filename, "r");
    if (!f) {
        perror("fopen public key");
        return NULL;
    }
    EVP_PKEY *pkey = PEM_read_PUBKEY(f, NULL, NULL, NULL);
    fclose(f);
    if (!pkey) {
        fprintf(stderr, "Failed to load public key from %s\n", filename);
        print_errors();
    }
    return pkey;
}

// 辅助：打印十六进制
void print_hex(const char *label, const unsigned char *data, size_t len) {
    printf("%s: ", label);
    for (size_t i = 0; i < len; i++) {
        printf("%02x", data[i]);
    }
    printf("\n");
}

// 主函数：演示全流程
int main(void) {
    const char *priv_file = "sm2_priv.pem";
    const char *pub_file = "sm2_pub.pem";
    const unsigned char msg[] = "Hello, SM2!";
    size_t msg_len = sizeof(msg) - 1; // 去掉 '\0'

    // 1. 生成密钥对
    if (generate_sm2_keypair(priv_file, pub_file) != 0) {
        return 1;
    }

    // 2. 加载密钥
    EVP_PKEY *priv_key = load_private_key(priv_file);
    EVP_PKEY *pub_key = load_public_key(pub_file);
    if (!priv_key || !pub_key) {
        EVP_PKEY_free(priv_key);
        EVP_PKEY_free(pub_key);
        return 1;
    }

    // 3. SM2 加密
    unsigned char *ciphertext = NULL;
    size_t ct_len = 0;
    if (sm2_encrypt(pub_key, msg, msg_len, &ciphertext, &ct_len) != 0) {
        fprintf(stderr, "SM2 encrypt failed\n");
        goto cleanup;
    }
    print_hex("Ciphertext", ciphertext, ct_len);

    // 4. SM2 解密
    unsigned char *plaintext = NULL;
    size_t pt_len = 0;
    if (sm2_decrypt(priv_key, ciphertext, ct_len, &plaintext, &pt_len) != 0) {
        fprintf(stderr, "SM2 decrypt failed\n");
        goto cleanup;
    }
    printf("Decrypted: %.*s\n", (int)pt_len, plaintext);

    // 5. SM2 签名
    unsigned char *signature = NULL;
    size_t sig_len = 0;
    if (sm2_sign(priv_key, msg, msg_len, &signature, &sig_len) != 0) {
        fprintf(stderr, "SM2 sign failed\n");
        goto cleanup;
    }
    print_hex("Signature", signature, sig_len);

    // 6. SM2 验签
    if (sm2_verify(pub_key, msg, msg_len, signature, sig_len) == 0) {
        printf("Signature verified successfully!\n");
    } else {
        fprintf(stderr, "Signature verification failed!\n");
        goto cleanup;
    }

cleanup:
    OPENSSL_free(ciphertext);
    OPENSSL_free(plaintext);
    OPENSSL_free(signature);
    EVP_PKEY_free(priv_key);
    EVP_PKEY_free(pub_key);
    return 0;
}
